"""
基于内存的简单向量数据库

这个模块提供了一个简单的基于内存的向量数据库实现，用于替代Milvus Lite
主要用于测试和演示，不需要额外的依赖
"""
import os
import numpy as np
from typing import List, Dict, Any, Optional


class SimpleVectorDB:
    """基于内存的简单向量数据库，用于存储和检索文档"""
    
    def __init__(self, collection_name: str = "documents", vector_dim: int = 1024):
        """
        初始化简单向量数据库
        
        Args:
            collection_name: 集合名称
            vector_dim: 向量维度
        """
        self.collection_name = collection_name
        self.vector_dim = vector_dim
        self.documents = []  # 存储文档和向量
        self.next_id = 0  # 用于生成唯一ID
    
    def add_documents(self, documents: List[Dict[str, Any]]) -> List[int]:
        """
        添加文档到向量数据库
        
        Args:
            documents: 文档列表，每个文档包含'id', 'vector', 'text'和'metadata'
        
        Returns:
            添加的文档ID列表
        """
        added_ids = []
        
        for doc in documents:
            # 使用提供的ID或生成新ID
            doc_id = doc.get('id', self.next_id)
            self.next_id = max(self.next_id, doc_id + 1)
            
            # 添加文档
            self.documents.append({
                'id': doc_id,
                'vector': np.array(doc['vector']),  # 转换为numpy数组以便计算
                'text': doc['text'],
                'metadata': doc.get('metadata', {})
            })
            
            added_ids.append(doc_id)
        
        print(f"添加了{len(documents)}个文档到向量数据库")
        return added_ids
    
    def search(self, query_vectors: List[List[float]], limit: int = 3, 
               filter_expr: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        在向量数据库中搜索相似文档
        
        Args:
            query_vectors: 查询向量列表
            limit: 返回结果的最大数量
            filter_expr: 过滤表达式（此简单实现暂不支持复杂过滤）
        
        Returns:
            搜索结果列表
        """
        if not self.documents:
            print("警告: 向量数据库为空")
            return []
        
        # 我们只处理第一个查询向量
        query_vector = np.array(query_vectors[0])
        
        # 计算余弦相似度
        results = []
        for doc in self.documents:
            # 计算余弦相似度
            similarity = self._cosine_similarity(query_vector, doc['vector'])
            
            # 如果有过滤表达式，这里应该进行过滤
            # 目前仅支持简单的元数据字段匹配
            if filter_expr and not self._simple_filter(doc['metadata'], filter_expr):
                continue
                
            results.append({
                'id': doc['id'],
                'distance': 1 - similarity,  # 转换相似度为距离
                'text': doc['text'],
                'metadata': doc['metadata']
            })
        
        # 按相似度排序（距离越小越相似）
        results.sort(key=lambda x: x['distance'])
        
        # 限制结果数量
        return results[:limit]
    
    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """计算两个向量的余弦相似度"""
        dot_product = np.dot(vec1, vec2)
        norm_a = np.linalg.norm(vec1)
        norm_b = np.linalg.norm(vec2)
        
        # 避免除以零
        if norm_a == 0 or norm_b == 0:
            return 0
            
        return dot_product / (norm_a * norm_b)
    
    def _simple_filter(self, metadata: Dict[str, Any], filter_expr: str) -> bool:
        """
        简单的元数据过滤器
        
        Args:
            metadata: 文档元数据
            filter_expr: 简单的过滤表达式，格式为"field=value"
        
        Returns:
            是否通过过滤器
        """
        try:
            if '=' not in filter_expr:
                return True
                
            field, value = filter_expr.split('=', 1)
            field = field.strip()
            value = value.strip().strip('"\'')  # 移除可能的引号
            
            # 检查元数据字段是否匹配
            return str(metadata.get(field)) == value
        except:
            # 过滤表达式错误，默认通过
            return True
    
    def delete_documents(self, ids: List[int] = None) -> List[int]:
        """
        从向量数据库中删除文档
        
        Args:
            ids: 要删除的文档ID列表
        
        Returns:
            已删除的文档ID列表
        """
        if not ids:
            return []
            
        deleted_ids = []
        remaining_docs = []
        
        for doc in self.documents:
            if doc['id'] in ids:
                deleted_ids.append(doc['id'])
            else:
                remaining_docs.append(doc)
        
        self.documents = remaining_docs
        return deleted_ids
    
    def drop_collection(self) -> None:
        """删除整个集合"""
        self.documents = []
        print(f"集合 {self.collection_name} 已清空")


# 全局实例，用于替代Milvus客户端
_simple_vector_db = None


def get_vector_db(db_path: str = None, collection_name: str = "documents", vector_dim: int = 1024):
    """获取向量数据库单例"""
    global _simple_vector_db
    if _simple_vector_db is None:
        _simple_vector_db = SimpleVectorDB(collection_name, vector_dim)
    return _simple_vector_db 